package com.xw.llm.test;

import io.vertx.core.AbstractVerticle;
import io.vertx.core.Vertx;
import io.vertx.core.buffer.Buffer;
import io.vertx.core.net.NetClient;
import io.vertx.core.net.NetServer;
import io.vertx.core.net.NetSocket;

public class MysqlProxy extends AbstractVerticle {

    // 目标MySQL服务器配置
    private static final String TARGET_HOST = System.getProperty("mysqlurl");
    private static final int TARGET_PORT = 3306;

    // 代理服务器配置
    private static final int PROXY_PORT = 3307;

    @Override
    public void start() {
        // 创建代理服务器
        NetServer proxyServer = vertx.createNetServer();

        // 创建连接到目标MySQL的客户端
        NetClient mysqlClient = vertx.createNetClient();

        proxyServer.connectHandler(clientSocket -> {
            System.out.println("客户端已连接");

            // 连接到目标MySQL服务器
            mysqlClient.connect(TARGET_PORT, TARGET_HOST, serverSocket -> {
                if (serverSocket.succeeded()) {
                    System.out.println("已连接到目标MySQL服务器");
                    NetSocket mysqlSocket = serverSocket.result();

                    // 处理客户端到服务器的数据流（拦截SQL）
                    clientSocket.handler(buffer -> {
                        // 尝试解析MySQL协议中的SQL语句
                        String sql = parseSqlFromBuffer(buffer);

                        if (sql != null) {
                            System.out.println("拦截到SQL: " + sql);

                            // 检查是否是需要拦截的特定SQL
//                            if (shouldIntercept(sql)) {
                            if (true) {
                                System.out.println("匹配到特定SQL，返回固定结果");
                                // 发送固定结果给客户端，而不转发到服务器
                                try {
                                    clientSocket.write(generateFixedResponse(sql));
                                } catch (Exception e) {
                                    throw new RuntimeException(e);
                                }
                                return;
                            }
                        }

                        // 非拦截SQL，转发到MySQL服务器
                        mysqlSocket.write(buffer);
                    });

                    // 处理服务器到客户端的数据流
                    mysqlSocket.handler(buffer -> {
                        // 直接转发服务器响应给客户端
                        try {
                            String string = buffer.getString(1, buffer.length());
                            System.out.println("服务端响应： " + buffer.getByte(0) + "   " + string);
                        } catch (Exception e) {
                            e.printStackTrace();
                        }
                        clientSocket.write(buffer);
                    });

                    // 处理连接关闭
                    clientSocket.closeHandler(v -> {
                        System.out.println("客户端连接关闭");
                        mysqlSocket.close();
                    });

                    mysqlSocket.closeHandler(v -> {
                        System.out.println("与MySQL服务器的连接关闭");
                        clientSocket.close();
                    });

                    // 处理异常
                    clientSocket.exceptionHandler(e -> {
                        System.err.println("客户端连接异常: " + e.getMessage());
                        e.printStackTrace();
                    });

                    mysqlSocket.exceptionHandler(e -> {
                        System.err.println("与MySQL服务器的连接异常: " + e.getMessage());
                        e.printStackTrace();
                    });
                } else {
                    System.err.println("无法连接到目标MySQL服务器: " + serverSocket.cause().getMessage());
                    clientSocket.close();
                }
            });
        });

        // 启动代理服务器
        proxyServer.listen(PROXY_PORT, res -> {
            if (res.succeeded()) {
                System.out.println("MySQL代理服务器已启动，监听端口: " + PROXY_PORT);
                System.out.println("转发到目标MySQL服务器: " + TARGET_HOST + ":" + TARGET_PORT);
            } else {
                System.err.println("启动代理服务器失败: " + res.cause().getMessage());
            }
        });
    }

    /**
     * 判断是否需要拦截该SQL
     */
    private boolean shouldIntercept(String sql) {
        // 这里可以根据需要修改拦截条件
        sql = sql.trim().toLowerCase();
        return sql.startsWith("select") && sql.contains("special_table");
    }

    /**
     * 解析MySQL协议中的SQL语句
     * 注意：这是简化实现，完整解析需要处理复杂的MySQL协议
     */
    private String parseSqlFromBuffer(Buffer buffer) {
        try {
            // MySQL客户端发送的命令包通常以命令类型开始
            // 0x03表示COM_QUERY命令，后面跟随SQL语句
//            if (buffer.length() > 0 && buffer.getByte(0) == 0x03) {
            if (buffer.length() > 0 && (buffer.getByte(0) == 53 || buffer.getByte(0) == 34)) {
                // 提取SQL语句（从第1个字节开始，UTF-8编码）
                return buffer.getString(1, buffer.length());
            }
        } catch (Exception e) {
            System.err.println("解析SQL失败: " + e.getMessage());
        }
        return null;
    }

    /**
     * 生成固定的响应数据
     * 注意：这是简化实现，实际需要返回符合MySQL协议的响应
     */
    private Buffer generateFixedResponse(String sql) throws Exception {
        Buffer response = Buffer.buffer();

// 1. 包头部（先预留，最后计算长度）
        int sequenceId = 0; // 包序号，从0开始
        byte[] header = new byte[4]; // 3字节长度 + 1字节序号
        response.appendBytes(header); // 先占位，最后填充

// 2. 列计数（3列，长度编码整数0x03）
        response.appendByte((byte) 0x03);

// 3. 列定义（以id列为例，简化版）
// 每个列定义包含：数据库名(空)、表名(空)、原始表名(空)、列名、类型等
        appendColumnDefinition(response, "id", (byte) 0x03); // 0x03表示INT类型
        appendColumnDefinition(response, "name", (byte) 0x0f); // 0x0f表示VARCHAR类型
        appendColumnDefinition(response, "time", (byte) 0x0f); // VARCHAR类型

// 4. 列定义结束EOF包（0xFE + 警告数0 + 服务器状态0）
        response.appendByte((byte) 0xFE);
        response.appendShortLE((short) 0); // 警告数
        response.appendShortLE((short) 0); // 服务器状态
        sequenceId++; // 序号递增

// 5. 行数据（1行，3个字段）
        response.appendByte((byte) 0x00); // 行开始（无NULL值）
// 字段1: id=1（长度1，值'1'）
        response.appendByte((byte) 0x01); // 长度
        response.appendByte((byte) '1');
// 字段2: name=test（长度4，值'test'）
        response.appendByte((byte) 0x04); // 长度
        response.appendBytes("test".getBytes());
// 字段3: time=2023-01-01（长度10，值'2023-01-01'）
        response.appendByte((byte) 0x0a); // 长度
        response.appendBytes("2023-01-01".getBytes());
        sequenceId++;

// 6. 行数据结束EOF包
        response.appendByte((byte) 0xFE);
        response.appendShortLE((short) 0);
        response.appendShortLE((short) 0);
        sequenceId++;

// 填充包头部（计算包体长度，更新序号）
        int bodyLength = response.length() - 4; // 总长度 - 头部4字节
        header[0] = (byte) (bodyLength & 0xFF);
        header[1] = (byte) ((bodyLength >> 8) & 0xFF);
        header[2] = (byte) ((bodyLength >> 16) & 0xFF);
        header[3] = (byte) sequenceId;
// 将头部写入buffer起始位置
        response.setBytes(0, header);


        return response;
    }

    private static void appendColumnDefinition(Buffer buf, String columnName, byte type) {
        // 数据库名（空，长度0）
        buf.appendByte((byte) 0x00);
        // 表名（空）
        buf.appendByte((byte) 0x00);
        // 原始表名（空）
        buf.appendByte((byte) 0x00);
        // 列名（长度编码：长度+值）
        buf.appendByte((byte) columnName.length()); // 列名长度
        buf.appendBytes(columnName.getBytes());
        // 原始列名（同列名）
        buf.appendByte((byte) columnName.length());
        buf.appendBytes(columnName.getBytes());
        // 字符集（utf8mb4 = 45）
        buf.appendShortLE((short) 45);
        // 列长度（INT=4，VARCHAR=255）
        buf.appendIntLE(type == 0x03 ? 4 : 255);
        // 列类型（如INT=0x03，VARCHAR=0x0f）
        buf.appendByte(type);
        // flags（默认0）
        buf.appendShortLE((short) 0);
    }

    public static void main(String[] args) {
        Vertx vertx = Vertx.vertx();
        vertx.deployVerticle(new MysqlProxy());
    }
}
